/* This file contains all pagefault-related routines */
#include "include.h"

typedef struct fault_handler_data {
    /* Initial args */
    sos_pcb *pcb;
    seL4_Word proc_vaddr;
    /* Intermediate state */
    page_table_entry *pte;
    sos_region *region;
    /* Information to resume caller */
    page_table_cb cb;
    void *cookie;
    int err;
} fault_handler_data;

static int _fault_handle_reply_task(fault_handler_data *data);
static void _cleanup(int err, void *cookie);
static void _fault_handle_have_pt(int err, void *cookie);
static void _map_it(seL4_Word sos_vaddr, void *_data);
static void _page_fault_reply(void *cookie);

static int _fault_handle_reply_task(fault_handler_data *data) {
    if (data->pcb->abort) {
        dprintf(6, "fault_handle_reply_task: aborting (PID %d zombie)\n",
                data->pcb->pid);
        data->err = SOS_PAGE_TABLE_ABORT;
    }

    // Reply to callback, free state, quit
    dprintf(6, "fault_handle_reply task: firing (%p)\n", data);
    (*(data->cb))(data->err, data->cookie);
    free(data);
    return SOS_TASK_SUCCESS;
}

static void _fault_handle_reply(fault_handler_data *data) {

    if (data->pcb->abort) {
        dprintf(6, "fault_handle_reply_task: aborting\n");
        data->err = SOS_PAGE_TABLE_ABORT;
    }

    int err;
    dprintf(7, "fault_handle_reply: scheduling reply task\n");
    err = sos_task_add_ready(&_fault_handle_reply_task, (void *) data,
                             SOS_TASK_PRIORITY_HIGH);
    if (err) {
        dprintf(1, "fault handler: couldn't alloc reply task (%d)\n", err);
        // Should manually tear down the process?
        // XXX Scratch pool
    }
    return;
}

/**
 *  The fault handler is asynchronous, since it calls frame_alloc between
 *  1 and 4 times in ordinary operation. All intermediate storage is done
 *  via a fault_handler_data struct on the heap. Reply to the caller happens
 *  via the task readyqueue, i.e. after this stackframe has been destroyed.
 *  The heap data is freed when replying to the caller.
 */
int sos_page_fault_handle(sos_pcb *pcb, seL4_Word proc_vaddr,
                          page_table_cb cb, void *cookie) {
    int err;

    if (proc_vaddr == 0) return SOS_PAGE_TABLE_SEGFAULT;

    // Grab some space to store async state for this operation
    fault_handler_data *data = malloc(sizeof(struct fault_handler_data));
    if (data == NULL) return SOS_PAGE_TABLE_OOM;
    bzero((void *) data, sizeof(struct fault_handler_data));

    // Copy the useful stuff into the heap
    data->pcb = pcb;
    data->proc_vaddr = proc_vaddr;
    data->cb = cb;
    data->cookie = cookie;
    data->err = SOS_PAGE_TABLE_SUCCESS;

    dprintf(6, "Handling a VM fault at %p! (PID: %d, PD: %p)\n",
            (void *) proc_vaddr, pcb->pid, pcb->sos_pd);

    // Sanity checks - has the process been initialised properly?
    if (pcb->sos_pd == 0) return SOS_PAGE_TABLE_INVAL;
    if (pcb->cspace == NULL) return SOS_PAGE_TABLE_INVAL;

    // Check VM region - figure out if this is supposed to succeed or not
    dprintf(6, "fault_handle: resolving VM region...\n");
    err = sos_resolve_region(&(data->pcb->regions), data->proc_vaddr, &(data->region));
    if (err != SOS_REGION_SUCCESS) {
        // No: call back with err = segfault, process gets torn down
        dprintf(1, "fault_handle: %p did not fall within any region for PID %d\n",
                proc_vaddr, pcb->pid);
        data->err = SOS_PAGE_TABLE_SEGFAULT;
        _fault_handle_reply(data);
        return SOS_PAGE_TABLE_SEGFAULT;
    } else {
        // Yes: valid access, try to map a frame in.
        // We need an entry in the page table to proceed; do a faulting walk
        dprintf(6, "fault_handle: walking the page table...\n");
        err = page_walk(data->pcb, data->proc_vaddr, data->region->attrs,
                        &_fault_handle_have_pt, (void *) data);
        if (err) {
            // XXX Unknown terrible error
            dprintf(5, "page_fault_handler: unknown terrible error occurred\n");
            data->err = SOS_PAGE_TABLE_SEGFAULT;
            _fault_handle_reply(data);
            return SOS_PAGE_TABLE_SEGFAULT;
        }
        return SOS_PAGE_TABLE_SUCCESS;
    }
}

/**
 * We've completed a faulting walk.
 * All levels of the SOS page table are in place.
 *
 * Four cases:
 *  1. Frame never mapped. We need a frame_alloc
 *  2. Frame unmapped but present. Call frame_referenced
 *  3. Frame mapped and present. Call frame_referenced
 *  4. Frame swapped out. Call frame_restore
 */
static void _fault_handle_have_pt(int err, void *cookie) {
    fault_handler_data *data = (fault_handler_data *) cookie;

    if (data->pcb->abort) {
        dprintf(6, "fault_handle_have_pt: aborting\n");
        data->err = SOS_PAGE_TABLE_ABORT;
        _fault_handle_reply(data);
        return;
    }

    dprintf(6, "page_fault_handle: page walk complete. have PTE\n");

    // Handle the error code from page_walk
    if (err) {
        // Something terrible has happened - region check succeeded, walk failed.
        dprintf(1, "fault_handle: Page walk failed unexpectedly (%d)\n", err);
        // The process needs to die. Caller will be watching for SEGFAULT.
        data->err = SOS_PAGE_TABLE_SEGFAULT;
        _fault_handle_reply(data);
        return;
    }

    // All good - grab PTE in constant time
    dprintf(8, "data->pcb: %p\tdata->proc_vaddr: %p\tpd: %p\n",
            data->pcb, data->proc_vaddr, data->pcb->sos_pd);
    err = page_walk_nofault(data->pcb, data->proc_vaddr, &(data->pte));
    dprintf(6, "page_fault_handle: page_walk_nofault returned ok\n");
    // If this fails, a core assumption is wrong and PTs have disappeared
    // (Or, alternatively, 'data' or the page table have been corrupted)
    if (err) {
        dprintf(6, "fault handler: unexpected page fault (%d)\n", err);
        data->err = err;
        _fault_handle_reply(data);
        return;
    }

    // Two possibilities: need fresh frame, or need old frame from swap.
    // Both branches need to perform the desired callback.
    if (data->pte->framenumber == 0) {
        // This page has never been accessed - need to map a frame
        dprintf(6, "page_fault_handle: need fresh frame for %p\n",
                data->proc_vaddr);
        err = page_fault(data->pcb, data->proc_vaddr, data->pte,
                         &_cleanup, (void *) data);
        if (err != SOS_PAGE_TABLE_SUCCESS) {
            // XXX Should do something sensible with scratch
            dprintf(2, "page_fault_handle: page_fault returned err %d\n", err);
            data->err = SOS_PAGE_TABLE_SEGFAULT;
            _fault_handle_reply(data);
            return;
        }
    } else if (data->pte->attrs.swapped) {
        // Need to grab the frame from swap
        err = frame_swapin(data->pte->framenumber, &_cleanup, (void *) data);
        if (err) {
            // Memory is trashed - segfault
            dprintf(6, "Something went wrong with a frame_swapin (%d)\n", err);
            data->err = err;
            _fault_handle_reply(data);
            return;
        }
    } else {
        dprintf(6, "fault_handle postwalk: frame has been chanced\n");

        // Not swapped, and has frame number. Frame is a swap candidate.
        // Call frame_reference, get it mapped back in.
        data->err = frame_referenced(data->pte->framenumber);

        data->err = SOS_PAGE_TABLE_SUCCESS;
        _fault_handle_reply(data);
        return;
    }

    return;
}

static void _cleanup(int err, void *cookie) {
    fault_handler_data *data = (fault_handler_data *) cookie;
    dprintf(5, "fault_handle_cleanup: running\n");
    data->err = err;
    _fault_handle_reply(data);
}


/**
 * We use a little more heap than necessary here - this page_fault function
 * can be called by other VM components.
 * XXX Should this be the case?
 */

/* Page fault - obtain a new frame and map it in */
/* Context: page table has been walked */
typedef struct page_fault_data {
    /* Initial arguments */
    sos_pcb *pcb;
    seL4_Word proc_vaddr;
    page_attrs attrs;
    page_table_entry *pte;
    /* Information to call back */
    page_table_cb cb;
    void *cookie;
    int err;
} page_fault_data;

static void _page_fault_reply(void *cookie) {
    dprintf(6, "page_fault_reply: firing\n");
    page_fault_data *data = (page_fault_data *) cookie;
    (*(data->cb))(data->err, data->cookie);
    free(data);
    return;
}

int page_fault(sos_pcb *pcb, seL4_Word proc_vaddr, page_table_entry *pte,
             page_table_cb cb, void *cookie) {
    if (pcb->abort) return SOS_PAGE_TABLE_OOM;

    dprintf(5, "page_fault: PID %d @ %p\n", pcb->pid, proc_vaddr);

    // Whack persistent data on the heap
    page_fault_data *data = malloc(sizeof(struct page_fault_data));
    if (data == NULL) {
        // XXX Do something intelligent with scratch
        dprintf(5, "page_fault: system is OOM\n");
        return SOS_PAGE_TABLE_OOM;
    }
    bzero((void *) data, sizeof(struct page_fault_data));


    // Stash useful arguments
    data->pcb = pcb;
    data->proc_vaddr = proc_vaddr;
    data->pte = pte;
    data->cb = cb;
    data->cookie = cookie;
    data->err = 0;

    // Set up frame_alloc
    frame_alloc(&_map_it, (void *) data);

    return SOS_PAGE_TABLE_SUCCESS;
}

/* This operation must be synchronous to avoid a race condition on grant / pin */
static void _map_it(seL4_Word sos_vaddr, void *_data) {
    page_fault_data *data = (page_fault_data *) _data;

    if (data->pcb->abort) {
        data->err = SOS_PAGE_TABLE_ABORT;
        _page_fault_reply(data);
        return SOS_PAGE_TABLE_ABORT;
    }

    int err;
    seL4_CPtr proc_cap = 0;

    // Check vaddr != null
    if (sos_vaddr == 0) {
        // No frames were available - we're out of memory.
        // Process is going to die. Pass the error upstream.
        dprintf(6, "Frame table is out of memory.\n");
        data->err = SOS_PAGE_TABLE_SEGFAULT;
        _page_fault_reply(data);
        return;
    }

    // Grant frame to the process
    dprintf(5, "Granting frame %p\n", (void*) sos_vaddr);
    err = frame_grant(sos_vaddr, data->pcb, data->proc_vaddr);
    if (err) {
        dprintf(6, "Failed to grant frame!\n");
        data->err = SOS_PAGE_TABLE_SEGFAULT;
        _page_fault_reply(data);
        return;
    }

    // Update the PTE with the framenumber
    dprintf(5, "Setting PTE framenumber...\n");
    data->pte->framenumber = frame_translate_vaddr(sos_vaddr);

    // Obtain cap
    err = frame_process_cap(sos_vaddr, data->pcb, &proc_cap);
    // Perform seL4 mapping
    assert(proc_cap);
    err = sos_page_map(data->pcb, data->proc_vaddr, proc_cap);
    data->err = err;

    dprintf(7, "page_map_it: page_map returned\n");
    // Perform success callback
    _page_fault_reply(data);

    return;
}
